"""
使用BaoStock同步股票日K数据脚本

本脚本用于从BaoStock获取股票日K数据并同步到数据库。
相关知识点：
- BaoStock数据获取
- 异步编程
- 批量数据处理
"""
import asyncio
import random
import sys
from datetime import date, datetime, timedelta
from typing import List, Tuple
from pathlib import Path

# 将项目根目录加入 sys.path，确保可导入 config 与 module_stock 包
PROJECT_ROOT = Path(__file__).resolve().parents[1]
if str(PROJECT_ROOT) not in sys.path:
    sys.path.insert(0, str(PROJECT_ROOT))
# 规避 config.env 的 argparse 仅允许 --env 的限制：先清空除脚本名以外的argv
sys.argv = [sys.argv[0]]

from config.database import AsyncSessionLocal
from module_stock.dao.basic_dao import BasicDao
from module_stock.service.baostock_daily_service import BaoStockDailyService
from utils.log_util import logger


async def _get_all_symbols(session, limit: int | None) -> List[str]:
    """
    获取所有股票代码
    
    :param session: 数据库会话
    :param limit: 限制数量，None表示获取全部
    :return: 股票代码列表
    """
    symbols = await BasicDao.get_all_symbols(session)
    if limit is not None and limit > 0:
        return symbols[:limit]
    return symbols


async def _sync_symbol_daily_data(
    session,
    symbol: str,
    start_date: str | None = None,
    end_date: str | None = None
) -> Tuple[str, int]:
    """
    同步某只股票的日K数据
    
    :param session: 数据库会话
    :param symbol: 股票代码
    :param start_date: 开始日期，格式 "YYYY-MM-DD"
    :param end_date: 结束日期，格式 "YYYY-MM-DD"
    :return: (股票代码, 同步的记录数)
    """
    try:
        # 先获取数据，用于统计数量
        records = await BaoStockDailyService.fetch_daily_data_from_baostock(
            symbol=symbol,
            start_date=start_date,
            end_date=end_date
        )
        
        record_count = len(records) if records else 0
        
        if record_count == 0:
            return symbol, 0
        
        # 同步到数据库
        await BaoStockDailyService.sync_daily_data_to_db(
            query_db=session,
            symbol=symbol,
            start_date=start_date,
            end_date=end_date
        )
        return symbol, record_count
    except Exception as e:
        logger.error(f"同步股票 {symbol} 失败: {e}")
        return symbol, 0


async def main(limit: int | None = None, sleep_min: float = 0.5, sleep_max: float = 1.0):
    """
    主函数：同步所有（或前N只）股票的日K数据
    
    :param limit: 限制同步的股票数量，None表示同步全部
    :param sleep_min: 最小休眠时间（秒）
    :param sleep_max: 最大休眠时间（秒）
    """
    async with AsyncSessionLocal() as session:
        # 获取所有股票代码
        symbols = await _get_all_symbols(session, limit)
        
        total_synced = 0
        processed = 0
        failed_symbols = []
        
        logger.info(f"开始同步 {len(symbols)} 只股票的日K数据")
        
        for symbol in symbols:
            try:
                # 同步数据
                sym, cnt = await _sync_symbol_daily_data(session, symbol)
                await session.commit()
                total_synced += cnt
                processed += 1
                
                if cnt > 0:
                    logger.info(f"[{processed}/{len(symbols)}] {symbol} 同步成功，新增 {cnt} 条")
                else:
                    logger.info(f"[{processed}/{len(symbols)}] {symbol} 没有新数据")
                    
            except Exception as e:
                await session.rollback()
                failed_symbols.append(symbol)
                logger.error(f"[{processed}/{len(symbols)}] {symbol} 同步失败: {e}")
                processed += 1
            finally:
                # 随机休眠，降低被风控概率
                if processed < len(symbols):
                    sleep_time = random.uniform(sleep_min, sleep_max)
                    await asyncio.sleep(sleep_time)
                
                # 每处理50只股票输出一次进度
                if processed % 50 == 0:
                    logger.info(f"进度: {processed}/{len(symbols)}，累计同步 {total_synced} 条数据")
        
        logger.info(f"同步完成！")
        logger.info(f"总计: {processed} 只股票，成功同步 {total_synced} 条数据")
        if failed_symbols:
            logger.warning(f"失败股票: {', '.join(failed_symbols)}")


async def sync_last_n_days(limit: int | None = None, days: int = 30, sleep_min: float = 0.5, sleep_max: float = 1.0):
    """
    同步最近N天的股票日K数据
    
    :param limit: 限制同步的股票数量，None表示同步全部
    :param days: 同步最近多少天的数据
    :param sleep_min: 最小休眠时间（秒）
    :param sleep_max: 最大休眠时间（秒）
    """
    end_date = date.today()
    start_date = end_date - timedelta(days=days)
    start_date_str = start_date.strftime("%Y-%m-%d")
    end_date_str = end_date.strftime("%Y-%m-%d")
    
    logger.info(f"同步最近 {days} 天的数据（{start_date_str} 至 {end_date_str}）")
    
    async with AsyncSessionLocal() as session:
        symbols = await _get_all_symbols(session, limit)
        
        total_synced = 0
        processed = 0
        failed_symbols = []
        
        for symbol in symbols:
            try:
                sym, cnt = await _sync_symbol_daily_data(
                    session,
                    symbol,
                    start_date=start_date_str,
                    end_date=end_date_str
                )
                await session.commit()
                total_synced += cnt
                processed += 1
                
                if cnt > 0:
                    logger.info(f"[{processed}/{len(symbols)}] {symbol} 同步成功，新增 {cnt} 条")
                    
            except Exception as e:
                await session.rollback()
                failed_symbols.append(symbol)
                logger.error(f"[{processed}/{len(symbols)}] {symbol} 同步失败: {e}")
                processed += 1
            finally:
                if processed < len(symbols):
                    sleep_time = random.uniform(sleep_min, sleep_max)
                    await asyncio.sleep(sleep_time)
                
                if processed % 50 == 0:
                    logger.info(f"进度: {processed}/{len(symbols)}，累计同步 {total_synced} 条数据")
        
        logger.info(f"同步完成！")
        logger.info(f"总计: {processed} 只股票，成功同步 {total_synced} 条数据")
        if failed_symbols:
            logger.warning(f"失败股票: {', '.join(failed_symbols)}")


if __name__ == "__main__":
    # 解析命令行参数
    args = sys.argv[1:]
    kv = {args[i].lstrip("-"): args[i + 1] for i in range(0, len(args) - 1, 2)} if len(args) >= 2 else {}
    
    limit = int(kv.get("limit")) if "limit" in kv else None
    sleep_min = float(kv.get("sleep-min")) if "sleep-min" in kv else 0.5
    sleep_max = float(kv.get("sleep-max")) if "sleep-max" in kv else 1.0
    days = int(kv.get("days")) if "days" in kv else None
    
    if days:
        # 同步最近N天的数据
        asyncio.run(sync_last_n_days(limit=limit, days=days, sleep_min=sleep_min, sleep_max=sleep_max))
    else:
        # 同步全部数据
        asyncio.run(main(limit=limit, sleep_min=sleep_min, sleep_max=sleep_max))

